package org.anno.chunkdb;

import com.google.appengine.api.datastore.Blob;
import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreServiceFactory;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyFactory;
import com.google.appengine.api.datastore.Query;
import com.google.common.base.Function;
import com.google.common.base.Predicates;
import com.google.common.collect.Iterables;
import com.google.protobuf.InvalidProtocolBufferException;

import java.io.IOException;

/**
 * Data access object for tables.
 */
public class TableDao {

  private static final Function<Entity, Metadata.Table> TABLE_DECODER =
      entity -> {
        Blob proto = (Blob) entity.getProperty("proto");
        try {
          return Metadata.Table.parseFrom(proto.getBytes());
        } catch (InvalidProtocolBufferException e) {
          e.printStackTrace();
          return null;
        }
      };

  private final DatastoreService store;

  public TableDao() {
    this(DatastoreServiceFactory.getDatastoreService());
  }

  public TableDao(DatastoreService store) {
    this.store = store;
  }

  public String create(Metadata.Table.Builder table) {
    Key key = store.allocateIds("Table", 1).getStart();
    String id = KeyFactory.keyToString(key);
    table.setTableId(id);
    Entity entity = new Entity(key);
    entity.setUnindexedProperty("proto", new Blob(table.build().toByteArray()));
    store.put(entity);
    return id;
  }

  public Metadata.Table get(String id) throws EntityNotFoundException, IOException {
    Key key = KeyFactory.stringToKey(id);
    Entity entity = store.get(key);
    Blob proto = (Blob) entity.getProperty("proto");
    return Metadata.Table.parseFrom(proto.getBytes());
  }

  public void update(Metadata.Table table) {
    Key key = KeyFactory.stringToKey(table.getTableId());
    Entity entity = new Entity(key);
    entity.setUnindexedProperty("proto", new Blob(table.toByteArray()));
    store.put(entity);
  }

  public void addChunk(String tableId, byte[] flatArray) {
    Key key = KeyFactory.stringToKey(tableId);
    Entity entity = new Entity("Chunk", key);
    entity.setUnindexedProperty("blob", new Blob(flatArray));
    store.put(entity);
  }

  public Iterable<Metadata.Table> listTables() {
    Query query = new Query("Table");
    return Iterables.filter(Iterables.transform(store.prepare(query).asIterable(), TABLE_DECODER),
        Predicates.notNull());
  }

  Iterable<byte[]> getChunks(String tableId) {
    Query query = new Query("Chunk", KeyFactory.stringToKey(tableId));
    return Iterables.transform(store.prepare(query).asIterable(), entity -> {
      Blob blob = (Blob) entity.getProperty("blob");
      return blob.getBytes();
    });
  }
}
